import pandas as pd
import numpy as np
import cv2
import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = ['KaiTi']
plt.rcParams['axes.unicode_minus'] = False
from tensorflow.keras.utils import to_categorical
from sklearn.utils import shuffle
import os




__Cannel_TASK_WAIT_HANDLE__=False

def Channel():
    __Cannel_TASK_WAIT_HANDLE__=True

def to_binary_string(x):
    if np.isnan(x):
        return np.nan
    return format(int(x), 'b')

def combine_columns_optimized(df, n):
    # 记录原始数据中是否存在空值
    original_has_na = df.isna().values.any()

    # 临时替换空值为特殊字符（这里使用 '_'），以确保它们在组合时不被忽略
    df.fillna('_', inplace=True)

    # 计算能整除n的列数
    num_combined_columns = df.shape[1] // n * n

    # 初始化一个新的numpy数组来存储组合后的列
    combined_columns = np.empty((df.shape[0], num_combined_columns // n), dtype=object)

    for i in range(0, num_combined_columns, n):
        if __Cannel_TASK_WAIT_HANDLE__ == False:
            # 使用numpy数组操作加速拼接过程
            combined_column = np.apply_along_axis(lambda x: ''.join(x.astype(str)), axis=1, arr=df.iloc[:, i:i + n].values)

            # 直接将组合后的列放入新的numpy数组
            combined_columns[:, i // n] = combined_column
        else:
            raise "用户已取消当前操作"

    # 将新的numpy数组转换为DataFrame，一次性添加所有新列
    new_column_names = [f'Column_{i}' for i in range(combined_columns.shape[1])]
    new_df = pd.DataFrame(combined_columns, columns=new_column_names, index=df.index)

    # 恢复原始空值标记
    if original_has_na:
        string = '_' * n
        new_df.replace(string, pd.NA, inplace=True)

    return new_df

def convert_to_decimal(x):
    if pd.isna(x):
        return np.nan  # 或者返回您希望的其他值，如0
    return int(x, base=2)

def load_AIS_data(n_train, n_test, kwn, unkwn, openrisk=0, num=0):

    data_train, data_test = balance_type(kwn, unkwn, n_train, n_test, openrisk)

    # num = 0  # 按2的num次方进行即为合并
    if __Cannel_TASK_WAIT_HANDLE__ == False:
        data_train['list'] = data_train['data'].apply(lambda x: [int(i) for i in x])
    else:
        raise "用户已取消当前操作"
    if __Cannel_TASK_WAIT_HANDLE__ == False:
        bin_train = pd.DataFrame(data_train['list'].tolist()).applymap(to_binary_string)
    else:
        raise "用户已取消当前操作"
    if __Cannel_TASK_WAIT_HANDLE__ == False:
        data_test['list'] = data_test['data'].apply(lambda x: [int(i) for i in x])
    else:
        raise "用户已取消当前操作"
    if __Cannel_TASK_WAIT_HANDLE__ == False:
        bin_test = pd.DataFrame(data_test['list'].tolist()).applymap(to_binary_string)
    else:
        raise "用户已取消当前操作"

    conbined_train = combine_columns_optimized(bin_train, 2 ** num)  # .replace(np.nan, 'NaN')
    conbined_test = combine_columns_optimized(bin_test, 2 ** num)  # .replace(np.nan, 'NaN')
    # 合并后转化为十进制
    bin_train = conbined_train.applymap(convert_to_decimal)
    bin_test = conbined_test.applymap(convert_to_decimal)

    packet_train = {
        'bin': [],
        'extractedfeature': [],
        'length': [],
        'label': [],
        'type': [],
        'img': [],
        'img_3': []
    }
    packet_train['bin'] = np.array(bin_train, dtype=np.int32)
    packet_train['length'] = bin_train.apply(lambda row: len(row.dropna()), axis=1)
    packet_train['label'] = data_train['label']
    packet_test = {
        'bin': [],
        'extractedfeature': [],
        'length': [],
        'label': [],
        'type': [],
        'img': [],
        'img_3': []
    }
    packet_test['bin'] = np.array(bin_test, dtype=np.int32)
    packet_test['length'] = bin_test.apply(lambda row: len(row.dropna()), axis=1)
    packet_test['label'] = data_test['label']
    bin_train = bin_train.fillna(int(0))
    bin_test = bin_test.fillna(int(0))
    packet_train['bin'] = np.array(bin_train, dtype=np.int32)
    packet_test['bin'] = np.array(bin_test, dtype=np.int32)

    # 重组成图像
    size = 16
    img_train = []
    img_test = []
    for g in range(len(packet_train['bin'])):
        #     img_train.append(cv2.resize(packet_train['bin'][g][:packet_train['length'][g]].astype('uint8'), (1, size * size)))
        img_train.append(cv2.resize(packet_train['bin'][g][:size * size].astype('uint8'), (1, size * size)))
    for i in range(len(img_train)):
        img_train[i] = img_train[i].reshape(size, size, 1)
    packet_train['img'] = img_train


    for g in range(len(packet_test['bin'])):
        #     img_test.append(cv2.resize(packet_test['bin'][g][:packet_test['length'][g]].astype('uint8'), (1, size * size)))
        img_test.append(cv2.resize(packet_test['bin'][g][:size * size].astype('uint8'), (1, size * size)))
    for i in range(len(img_test)):
        img_test[i] = img_test[i].reshape(size, size, 1)
    packet_test['img'] = img_test

    # 假设您想要按照字母顺序编码类别标签
    custom_mapping = {}
    for k, label in enumerate(kwn + unkwn):
        custom_mapping[str(label)] = k
    combined_label = pd.DataFrame(
        np.concatenate((packet_train['label'].values, packet_test['label'].values), axis=0)).astype(str)
    # 使用 replace() 函数将原始类别标签替换为自定义编码
    combined_labels_mapped = combined_label.replace(custom_mapping)
    combined_labels = pd.DataFrame(to_categorical(combined_labels_mapped))

    # 分离训练集和测试集的编码结果
    y_train_onehot = np.array(combined_labels.iloc[:len(packet_train['length']), :len(kwn)])
    y_test_onehot = np.array(combined_labels.iloc[len(packet_train['length']):, :])
    y_train = y_train_onehot.argmax(axis=-1)
    y_test = y_test_onehot.argmax(axis=-1)

    packet_train['type'] = y_train
    packet_test['type'] = y_test

    return packet_train, packet_test

def load_WAIBU_train(data_train, kwn, num=0):

    # num = 0  # 按2的num次方进行即为合并
    data_train['list'] = data_train['data'].apply(lambda x: [int(i) for i in x[:256]])
    bin_train = pd.DataFrame(data_train['list'].tolist()).applymap(to_binary_string)

    conbined_train = combine_columns_optimized(bin_train, 2 ** num)  # .replace(np.nan, 'NaN')
    # 合并后转化为十进制
    bin_train = conbined_train.applymap(convert_to_decimal)

    packet_train = {
        'bin': [],
        'extractedfeature': [],
        'length': [],
        'label': [],
        'type': [],
        'img': [],
        'img_3': []
    }

    packet_train['length'] = bin_train.apply(lambda row: len(row.dropna()), axis=1)
    packet_train['label'] = data_train['label']
    bin_train = bin_train.fillna(int(0))
    packet_train['bin'] = np.array(bin_train, dtype=np.int32)

    # 重组成图像
    size = 16
    img_train = []
    for g in range(len(packet_train['bin'])):
        #     img_train.append(cv2.resize(packet_train['bin'][g][:packet_train['length'][g]].astype('uint8'), (1, size * size)))
        img_train.append(cv2.resize(packet_train['bin'][g][:size * size].astype('uint8'), (1, size * size)))
    for i in range(len(img_train)):
        img_train[i] = img_train[i].reshape(size, size, 1)
    packet_train['img'] = img_train

    # 假设您想要按照字母顺序编码类别标签
    custom_mapping = {}
    for k, label in enumerate(kwn):
        custom_mapping[str(label)] = k
    train_label = pd.DataFrame((packet_train['label'].values)).astype(str)
    # 使用 replace() 函数将原始类别标签替换为自定义编码
    train_labels_mapped = train_label.replace(custom_mapping)
    train_labels = pd.DataFrame(to_categorical(train_labels_mapped))

    # 分离训练集和测试集的编码结果
    y_train_onehot = np.array(train_labels)
    y_train = y_train_onehot.argmax(axis=-1)

    packet_train['type'] = y_train

    return packet_train

def load_WAIBU_test(data_test, num=0):

    # num = 0  # 按2的num次方进行即为合并
    data_test['list'] = data_test['data'].apply(lambda x: [int(i) for i in x[:256]])
    bin_test = pd.DataFrame(data_test['list'].tolist()).applymap(to_binary_string)

    conbined_test = combine_columns_optimized(bin_test, 2 ** num)  # .replace(np.nan, 'NaN')
    # 合并后转化为十进制
    bin_test = conbined_test.applymap(convert_to_decimal)
    packet_test = {
        'bin': [],
        'extractedfeature': [],
        'length': [],
        'label': [],
        'type': [],
        'img': [],
        'img_3': []
    }
    packet_test['length'] = bin_test.apply(lambda row: len(row.dropna()), axis=1)
    packet_test['label'] = data_test['label']
    bin_test = bin_test.fillna(int(0))
    packet_test['bin'] = np.array(bin_test, dtype=np.int32)

    # 重组成图像
    size = 16
    img_test = []

    for g in range(len(packet_test['bin'])):
        #     img_test.append(cv2.resize(packet_test['bin'][g][:packet_test['length'][g]].astype('uint8'), (1, size * size)))
        img_test.append(cv2.resize(packet_test['bin'][g][:size * size].astype('uint8'), (1, size * size)))
    for i in range(len(img_test)):
        img_test[i] = img_test[i].reshape(size, size, 1)
    packet_test['img'] = img_test

    return packet_test